#include "mini_chibicc.h"
#include <errno.h>
#include <limits.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <ctype.h>
#include <stdlib.h>

#ifdef _WIN32
#  include <direct.h>
#  ifndef PATH_MAX
#    define PATH_MAX _MAX_PATH
#  endif
#else
#  include <sys/stat.h>
#  include <unistd.h>
#endif

// Global symbols exposed elsewhere in the codebase that still expect
// these definitions to exist even though the driver logic has changed.
StringArray include_paths = {};
bool opt_fcommon = true;
bool opt_fpic = false;
char *base_file = NULL;

typedef struct {
  StringArray input_files;
  StringArray include_dirs;
  StringArray defines;
  StringArray undefines;
  StringArray link_libs;
  StringArray link_search_paths;
  char *output_dir;
  char *crate_name;
} Options;

static void usage(void) {
  fprintf(stderr, "Usage: translate_chibicc [options] <input.c>...\n");
  fprintf(stderr, "  -o <path>           Output directory for the Cargo project (default: out)\n");
  fprintf(stderr, "  -I <dir>            Add include directory (repeatable)\n");
  fprintf(stderr, "  -D <name[=value]>   Define macro for the C build\n");
  fprintf(stderr, "  -U <name>           Undefine macro for the C build\n");
  fprintf(stderr, "  -l <name>           Link with an additional native library\n");
  fprintf(stderr, "  -L <dir>            Add native library search path\n");
  fprintf(stderr, "  --crate-name <name> Override generated crate name\n");
  exit(1);
}

static char *xstrdup(const char *s) {
  size_t n = strlen(s) + 1;
  char *dup = (char *)malloc(n);
  if (!dup)
    error("malloc failed in xstrdup: %s", strerror(errno));
  memcpy(dup, s, n);
  return dup;
}

static char *xstrndup(const char *s, size_t n) {
  size_t len = 0;
  while (len < n && s[len])
    len++;
  char *dup = (char *)malloc(len + 1);
  if (!dup)
    error("malloc failed in xstrndup: %s", strerror(errno));
  if (len)
    memcpy(dup, s, len);
  dup[len] = '\0';
  return dup;
}

static void ensure_dir_exists(const char *path) {
#ifdef _WIN32
  if (_mkdir(path) == -1 && errno != EEXIST)
    error("failed to create directory %s: %s", path, strerror(errno));
#else
  if (mkdir(path, 0777) == -1 && errno != EEXIST)
    error("failed to create directory %s: %s", path, strerror(errno));
#endif
}

static void make_dir_all(const char *path) {
  if (!path || !*path)
    return;

  char *tmp = xstrdup(path);
  size_t len = strlen(tmp);

  while (len > 1 && (tmp[len - 1] == '/' || tmp[len - 1] == '\\')) {
    tmp[len - 1] = '\0';
    len--;
  }

  for (size_t i = 0; tmp[i]; i++) {
    if (tmp[i] != '/' && tmp[i] != '\\')
      continue;
    if (i == 0)
      continue;
    if (i == 2 && tmp[1] == ':')
      continue;

    char saved = tmp[i];
    tmp[i] = '\0';
    if (*tmp)
      ensure_dir_exists(tmp);
    tmp[i] = saved;

    while (tmp[i + 1] == '/' || tmp[i + 1] == '\\')
      i++;
  }

  if (*tmp)
    ensure_dir_exists(tmp);
  free(tmp);
}

static char *absolute_path(const char *path) {
  if (!path)
    return NULL;
#ifdef _WIN32
  char buf[PATH_MAX];
  if (!_fullpath(buf, path, PATH_MAX))
    error("failed to canonicalize path: %s", path);
  return xstrdup(buf);
#else
  char *resolved = realpath(path, NULL);
  if (!resolved)
    error("failed to canonicalize path: %s: %s", path, strerror(errno));
  return resolved;
#endif
}

static void slash_to_forward(char *path) {
  if (!path)
    return;
  for (char *p = path; *p; p++)
    if (*p == '\\')
      *p = '/';
}

static char *escape_rust_string(const char *s) {
  size_t len = strlen(s);
  char *buf = calloc(1, len * 4 + 3);
  if (!buf)
    error("calloc failed: %s", strerror(errno));

  char *p = buf;
  *p++ = '"';
  for (size_t i = 0; i < len; i++) {
    unsigned char c = (unsigned char)s[i];
    switch (c) {
    case '\\': *p++ = '\\'; *p++ = '\\'; break;
    case '"': *p++ = '\\'; *p++ = '"'; break;
    case '\n': *p++ = '\\'; *p++ = 'n'; break;
    case '\r': *p++ = '\\'; *p++ = 'r'; break;
    case '\t': *p++ = '\\'; *p++ = 't'; break;
    default:
      if (c < 0x20) {
        sprintf(p, "\\x%02x", c);
        p += 4;
      } else {
        *p++ = c;
      }
    }
  }
  *p++ = '"';
  *p = '\0';
  return buf;
}

static char *sanitize_crate_name(const char *name) {
  if (!name || !*name)
    return xstrdup("translated_project");

  size_t len = strlen(name);
  char *buf = calloc(1, len + 8);
  if (!buf)
    error("calloc failed: %s", strerror(errno));

  size_t j = 0;
  for (size_t i = 0; i < len; i++) {
    unsigned char c = (unsigned char)name[i];
    if (isalnum(c) || c == '_')
      buf[j++] = tolower(c);
    else
      buf[j++] = '_';
  }

  if (j == 0)
    buf[j++] = 't';
  buf[j] = '\0';

  if (!(isalpha((unsigned char)buf[0]) || buf[0] == '_')) {
    char *prefixed = format("c%s", buf);
    free(buf);
    buf = prefixed;
  }

  return buf;
}

static char *derive_crate_name(const char *path) {
  char *tmp = xstrdup(path);
  char *base = basename(tmp);
  char *dot = strrchr(base, '.');
  if (dot)
    *dot = '\0';
  char *result = sanitize_crate_name(base);
  free(tmp);
  return result;
}

static void parse_args(int argc, char **argv, Options *opt) {
  bool parsing_opts = true;

  for (int i = 1; i < argc; i++) {
    char *arg = argv[i];

    if (parsing_opts && strcmp(arg, "--") == 0) {
      parsing_opts = false;
      continue;
    }

    if (parsing_opts && arg[0] == '-') {
      if (!strcmp(arg, "-h") || !strcmp(arg, "--help"))
        usage();

      if (!strcmp(arg, "-o")) {
        if (++i >= argc)
          usage();
        opt->output_dir = xstrdup(argv[i]);
        continue;
      }

      if (!strncmp(arg, "-o", 2) && arg[2]) {
        opt->output_dir = xstrdup(arg + 2);
        continue;
      }

      if (!strcmp(arg, "--crate-name")) {
        if (++i >= argc)
          usage();
        opt->crate_name = sanitize_crate_name(argv[i]);
        continue;
      }

      if (!strncmp(arg, "-I", 2)) {
        char *val = arg + 2;
        if (!*val) {
          if (++i >= argc)
            usage();
          val = argv[i];
        }
        strarray_push(&opt->include_dirs, xstrdup(val));
        continue;
      }

      if (!strncmp(arg, "-D", 2)) {
        char *val = arg + 2;
        if (!*val) {
          if (++i >= argc)
            usage();
          val = argv[i];
        }
        strarray_push(&opt->defines, xstrdup(val));
        continue;
      }

      if (!strncmp(arg, "-U", 2)) {
        char *val = arg + 2;
        if (!*val) {
          if (++i >= argc)
            usage();
          val = argv[i];
        }
        strarray_push(&opt->undefines, xstrdup(val));
        continue;
      }

      if (!strcmp(arg, "-l")) {
        if (++i >= argc)
          usage();
        strarray_push(&opt->link_libs, xstrdup(argv[i]));
        continue;
      }

      if (!strncmp(arg, "-l", 2) && arg[2]) {
        strarray_push(&opt->link_libs, xstrdup(arg + 2));
        continue;
      }

      if (!strcmp(arg, "-L")) {
        if (++i >= argc)
          usage();
        strarray_push(&opt->link_search_paths, xstrdup(argv[i]));
        continue;
      }

      if (!strncmp(arg, "-L", 2) && arg[2]) {
        strarray_push(&opt->link_search_paths, xstrdup(arg + 2));
        continue;
      }

      if (!strcmp(arg, "-c") || !strcmp(arg, "-S") || !strcmp(arg, "-E") ||
          !strcmp(arg, "-M") || !strcmp(arg, "-MD") || !strcmp(arg, "-MMD") ||
          !strcmp(arg, "-MP") || !strcmp(arg, "-MF") || !strcmp(arg, "-MT") ||
          !strcmp(arg, "-pthread") || !strcmp(arg, "-fno-common") ||
          !strcmp(arg, "-fno-builtin") || !strcmp(arg, "-fno-strict-aliasing") ||
          !strcmp(arg, "-fno-stack-protector") || !strcmp(arg, "-fno-omit-frame-pointer") ||
          !strcmp(arg, "-pie") || !strcmp(arg, "-static") || !strcmp(arg, "-shared") ||
          !strncmp(arg, "-O", 2) || !strncmp(arg, "-W", 2) || !strncmp(arg, "-g", 2) ||
          !strncmp(arg, "-std=", 5) || !strncmp(arg, "-m", 2)) {
        continue;
      }

      error("unsupported option: %s", arg);
    } else {
      strarray_push(&opt->input_files, xstrdup(arg));
    }
  }

  if (opt->input_files.len == 0)
    usage();

  if (!opt->output_dir)
    opt->output_dir = xstrdup("out");

  if (!opt->crate_name)
    opt->crate_name = derive_crate_name(opt->input_files.data[0]);
}

static FILE *xfopen(const char *path, const char *mode) {
  FILE *fp = fopen(path, mode);
  if (!fp)
    error("cannot open %s: %s", path, strerror(errno));
  return fp;
}

static void write_text_file(const char *path, const char *content) {
  FILE *fp = xfopen(path, "wb");
  size_t len = strlen(content);
  if (len && fwrite(content, 1, len, fp) != len) {
    int err = errno;
    fclose(fp);
    error("failed to write %s: %s", path, strerror(err));
  }
  fclose(fp);
}

static void emit_cargo_toml(const Options *opt) {
  char *path = format("%s/Cargo.toml", opt->output_dir);
  char *content = format(
      "[package]\n"
      "name = \"%s\"\n"
      "version = \"0.1.0\"\n"
      "edition = \"2021\"\n"
      "publish = false\n\n"
      "[dependencies]\n\n"
      "[build-dependencies]\ncc = \"1.0\"\n\n"
      "[profile.dev]\npanic = \"abort\"\n\n"
      "[profile.release]\npanic = \"abort\"\n",
      opt->crate_name);

  write_text_file(path, content);
  free(path);
  free(content);
}

static void emit_build_rs(const Options *opt, const StringArray *abs_inputs,
                          const StringArray *abs_includes,
                          const StringArray *abs_link_search) {
  char *path = format("%s/build.rs", opt->output_dir);
  FILE *fp = xfopen(path, "wb");
  free(path);

  fprintf(fp, "fn main() {\n");
  fprintf(fp, "    let mut build = cc::Build::new();\n");
  fprintf(fp, "    build.warnings(false);\n");
  fprintf(fp, "    build.flag_if_supported(\"-std=c11\");\n");
  fprintf(fp, "    build.define(\"main\", Some(\"translated_main\"));\n");

  for (int i = 0; i < abs_includes->len; i++) {
    char *lit = escape_rust_string(abs_includes->data[i]);
    fprintf(fp, "    build.include(%s);\n", lit);
    fprintf(fp, "    println!(\"cargo:rerun-if-changed={}\", %s);\n", lit);
    free(lit);
  }

  for (int i = 0; i < opt->defines.len; i++) {
    char *def = opt->defines.data[i];
    char *eq = strchr(def, '=');
    if (eq) {
  char *name = xstrndup(def, (size_t)(eq - def));
      char *value = escape_rust_string(eq + 1);
      char *name_lit = escape_rust_string(name);
      fprintf(fp, "    build.define(%s, Some(%s));\n", name_lit, value);
      free(name);
      free(value);
      free(name_lit);
    } else {
      char *name_lit = escape_rust_string(def);
      fprintf(fp, "    build.define(%s, None);\n", name_lit);
      free(name_lit);
    }
  }

  for (int i = 0; i < opt->undefines.len; i++) {
    char *name_lit = escape_rust_string(opt->undefines.data[i]);
    fprintf(fp, "    build.undefine(%s);\n", name_lit);
    free(name_lit);
  }

  for (int i = 0; i < abs_inputs->len; i++) {
    char *lit = escape_rust_string(abs_inputs->data[i]);
    fprintf(fp, "    build.file(%s);\n", lit);
    fprintf(fp, "    println!(\"cargo:rerun-if-changed={}\", %s);\n", lit);
    free(lit);
  }

  for (int i = 0; i < abs_link_search->len; i++) {
    char *lit = escape_rust_string(abs_link_search->data[i]);
    fprintf(fp, "    println!(\"cargo:rustc-link-search=native={}\", %s);\n", lit);
    free(lit);
  }

  for (int i = 0; i < opt->link_libs.len; i++) {
    char *lit = escape_rust_string(opt->link_libs.data[i]);
    fprintf(fp, "    println!(\"cargo:rustc-link-lib={}\", %s);\n", lit);
    free(lit);
  }

  char *crate_lit = escape_rust_string(opt->crate_name);
  fprintf(fp, "    build.compile(%s);\n", crate_lit);
  free(crate_lit);

  fprintf(fp, "}\n");
  fclose(fp);
}

static void emit_main_rs(const Options *opt) {
  char *src_dir = format("%s/src", opt->output_dir);
  make_dir_all(src_dir);
  free(src_dir);

  char *path = format("%s/src/main.rs", opt->output_dir);
  const char *content =
      "use std::ffi::CString;\n"
      "use std::os::raw::c_char;\n"
      "\n"
      "extern \"C\" {\n"
      "    fn translated_main(argc: i32, argv: *mut *mut c_char) -> i32;\n"
      "}\n"
      "\n"
      "fn main() {\n"
      "    let args: Vec<CString> = std::env::args()\n"
      "        .map(|arg| CString::new(arg).expect(\"argument contains null byte\"))\n"
      "        .collect();\n"
      "    let mut ptrs: Vec<*mut c_char> = args\n"
      "        .iter()\n"
      "        .map(|arg| arg.as_ptr() as *mut c_char)\n"
      "        .collect();\n"
      "    ptrs.push(std::ptr::null_mut());\n"
      "    let exit_code = unsafe { translated_main(args.len() as i32, ptrs.as_mut_ptr()) };\n"
      "    std::process::exit(exit_code);\n"
      "}\n";

  write_text_file(path, content);
  free(path);
}

static void emit_project(Options *opt) {
  make_dir_all(opt->output_dir);

  StringArray abs_inputs = {};
  for (int i = 0; i < opt->input_files.len; i++) {
    char *abs = absolute_path(opt->input_files.data[i]);
    slash_to_forward(abs);
    strarray_push(&abs_inputs, abs);
  }

  StringArray abs_includes = {};
  for (int i = 0; i < opt->include_dirs.len; i++) {
    char *abs = absolute_path(opt->include_dirs.data[i]);
    slash_to_forward(abs);
    strarray_push(&abs_includes, abs);
  }

  StringArray abs_link_search = {};
  for (int i = 0; i < opt->link_search_paths.len; i++) {
    char *abs = absolute_path(opt->link_search_paths.data[i]);
    slash_to_forward(abs);
    strarray_push(&abs_link_search, abs);
  }

  emit_cargo_toml(opt);
  emit_main_rs(opt);
  emit_build_rs(opt, &abs_inputs, &abs_includes, &abs_link_search);

  fprintf(stderr, "Generated Cargo project at %s\n", opt->output_dir);
}

int main(int argc, char **argv) {
  Options opt = {};
  parse_args(argc, argv, &opt);
  emit_project(&opt);
  return 0;
}
